import cv2
import numpy as np
import math
from cam_calibration import (
    CamCalibration,
    COLS,
    ROWS,
    SQUARE_SIZE,
    JSON_FILE,
    IMG_PATH,
    JSON_PATH,
)
import os
import json
import time
from camera import Camera, ESC_KEY


class CheckerboardPoseEstimation:
    def __init__(self, intrinsic_params):
        self.mtx = np.array(intrinsic_params["mtx"])
        self.dist = np.array(intrinsic_params["dist"])
        self.obj_points = self._get_object_points()
        self.chessboard_size = (COLS, ROWS)

    def _get_object_points(self):
        objp = np.zeros((ROWS * COLS, 3), np.float32)
        objp[:, :2] = np.mgrid[0:COLS, 0:ROWS].T.reshape(-1, 2)
        objp = SQUARE_SIZE * objp
        return objp

    # 帅：返回的结果的单位是mm和度数，并且这里的pitch, yaw, roll是根据博客中来的，所以可能与飞控的不太一样，注意一下
    def estimate_pose(self, frame):
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(gray, self.chessboard_size, None)
        if ret:
            cv2.drawChessboardCorners(frame, self.chessboard_size, corners, ret)
            _, rvec, tvec = cv2.solvePnP(self.obj_points, corners, self.mtx, self.dist)
            rvec_matrix = cv2.Rodrigues(rvec)[0]
            proj_matrix = np.hstack((rvec_matrix, tvec))
            eulerAngles = cv2.decomposeProjectionMatrix(proj_matrix)[6]
            x, y, z = map(float, tvec)
            distance = math.sqrt(x**2 + y**2 + z**2)
            pitch, yaw, roll = map(float, eulerAngles)
            return ret, (x, y, z, pitch, yaw, roll), distance
        else:
            return ret, (0, 0, 0, 0, 0, 0), 0


if __name__ == "__main__":
    camera = Camera()
    if not camera.configure_pipeline():
        exit()
    if not camera.start_pipeline():
        exit()

    if JSON_FILE:
        json_file_path = os.path.join(
            JSON_PATH, "intrinsic_params.json"
        )  # JSON文件的路径
        if os.path.exists(json_file_path):
            with open(json_file_path, "r") as f:
                intrinsic = json.load(f)
        else:
            print(f"JSON file {json_file_path} not found!")
            exit()
    else:
        calib = CamCalibration()
        intrinsic = calib.calibrate_camera(IMG_PATH)

    pose_estimator = CheckerboardPoseEstimation(intrinsic)

    start_time = time.time()
    frame_count = 0
    try:
        while True:
            color_image = camera.get_frame()
            if color_image is None:
                continue
            ret, cam_pose, distance = pose_estimator.estimate_pose(color_image)
            if ret:
                cv2.putText(
                    color_image,
                    f"x: {cam_pose[0]:.2f}mm, y: {cam_pose[1]:.2f}mm, z: {cam_pose[2]:.2f}mm",
                    (10, color_image.shape[0] - 50),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.7,
                    (0, 255, 0),
                    2,
                )
                cv2.putText(
                    color_image,
                    f"dist: {distance:.2f}mm, pitch: {cam_pose[3]:.2f}, yaw: {cam_pose[4]:.2f}, roll: {cam_pose[5]:.2f}",
                    (10, color_image.shape[0] - 20),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.7,
                    (0, 255, 0),
                    2,
                )
            else:
                cv2.putText(
                    color_image,
                    "Unable to Detect Chessboard",
                    (20, color_image.shape[0] - 20),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    1.3,
                    (0, 0, 255),
                    3,
                )
            cv2.imshow("Color Viewer", color_image)

            frame_count += 1
            elapsed_time = time.time() - start_time
            if elapsed_time > 1.0:
                fps = frame_count / elapsed_time
                print(f"FPS: {fps:.2f}")
                start_time = time.time()
                frame_count = 0

            key = cv2.waitKey(1)
            if key == ord("q") or key == ESC_KEY:
                break
    except KeyboardInterrupt:
        pass
    finally:
        camera.stop_pipeline()
        cv2.destroyAllWindows()
